Add memory projection cli #273
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
yuankaichen-amd
requested review from
Xiaoming-AMD,
limou102 and
wenxie-amd
as code owners
wenxie-amd
force-pushed
the
dev/yk_perf_model
branch
from
89ba376 to
d6d0436
Compare
yuankaichen-amd
force-pushed
the
dev/yk_perf_model
branch
from
d6d0436 to
10c106b
Compare
araina-amd
reviewed
Nov 10, 2025
| total += num_tokens * self.config.model_config.ffn_hidden_size * 2 # bf16 | ||
| # Second Gemm | ||
| total += num_tokens * self.config.model_config.ffn_hidden_size * 2 # bf16 | ||
| return total |
Contributor
There was a problem hiding this comment.
non-swiglu and tensor model parallelism is not taken into account here.
wenxie-amd
force-pushed
the
dev/yk_perf_model
branch
from
17fd561 to
05da689
Compare
Xiaoming-AMD
approved these changes
Nov 12, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This cli uses *-pretrain.yaml as input and projects memory usage on a worker (default rank=0). It also prints a breakdown of parameter numbers and activation memory usage for each submodule from the model.
Currently it only supports Megatron config files.
Example usage:
NNODES=96 PRIMUS_MODEL=deepseek_proxy_2T PRIMUS_EP=16 PRIMUS_PP=12 bash runner/primus-cli direct --no-gpu --single -- projection memory --config examples/megatron/configs/MI300X/deepseek_v2-pretrain.yaml
Example output:
Total Number of Parameters: 2049.574961 Billion (2,049,574,961,152.0)
[embedding]
Params: 0.822084 Billion (822,083,584)
Activation Memory: 0.0625 GB
[dense_transformer_layer]
Params: 0.453018 Billion (453,017,600.0)
Activation Memory: 0.6250 GB
...........
...........
...........
[Primus:Projection] Memory Projection Summary on Rank 0:
Params: 13.006799 Billion (13,006,798,848.0)
Param+Optimizer Memory: 78.7379 GB
Activation Memory (per batch size 1, seq len 4096): 395.3125 GB
Projected Total Memory: 474.0504 GB
Follow-up work:
(1) consolidate parameters in projection config;
(2) Torchtitan support.